'use client';

import { useControllableState } from '@radix-ui/react-use-controllable-state';
import {
    Collapsible,
    CollapsibleContent,
    CollapsibleTrigger
} from '@/components/ui/collapsible';
import { cn } from '@/lib/utils';
import { BrainIcon, ChevronDownIcon } from 'lucide-react';
import type { ComponentProps } from 'react';
import { createContext, memo, useContext, useEffect, useState } from 'react';
import { Response } from './response';

type ReasoningContextValue = {
    isStreaming: boolean;
    isOpen: boolean;
    setIsOpen: (open: boolean) => void;
    duration: number;
};

const ReasoningContext = createContext<ReasoningContextValue | null>(null);

const useReasoning = () => {
    const context = useContext(ReasoningContext);
    if (!context) {
        throw new Error('Reasoning components must be used within Reasoning');
    }
    return context;
};

export type ReasoningProps = ComponentProps<typeof Collapsible> & {
    isStreaming?: boolean;
    open?: boolean;
    defaultOpen?: boolean;
    onOpenChange?: (open: boolean) => void;
    duration?: number;
};

const AUTO_CLOSE_DELAY = 1000;

export const Reasoning = memo(
    ({
         className,
         isStreaming = false,
         open,
         defaultOpen = false,
         onOpenChange,
         duration: durationProp,
         children,
         ...props
     }: ReasoningProps) => {
        const [isOpen, setIsOpen] = useControllableState({
            prop: open,
            defaultProp: defaultOpen,
            onChange: onOpenChange
        });
        const [duration, setDuration] = useControllableState({
            prop: durationProp,
            defaultProp: 0
        });

        const [hasAutoClosedRef, setHasAutoClosedRef] = useState(false);
        const [startTime, setStartTime] = useState<number | null>(null);

        // Track duration when streaming starts and ends
        useEffect(() => {
            if (isStreaming) {
                if (startTime === null) {
                    setStartTime(Date.now());
                }
            } else if (startTime !== null) {
                setDuration(Math.round((Date.now() - startTime) / 1000));
                setStartTime(null);
            }
        }, [isStreaming, startTime, setDuration]);

        // Auto-open when streaming starts, auto-close when streaming ends (once only)
        useEffect(() => {
            if (isStreaming && !isOpen) {
                setIsOpen(true);
            } else if (!isStreaming && isOpen && !defaultOpen && !hasAutoClosedRef) {
                // Add a small delay before closing to allow user to see the content
                const timer = setTimeout(() => {
                    setIsOpen(false);
                    setHasAutoClosedRef(true);
                }, AUTO_CLOSE_DELAY);
                return () => clearTimeout(timer);
            }
        }, [isStreaming, isOpen, defaultOpen, setIsOpen, hasAutoClosedRef]);

        const handleOpenChange = (newOpen: boolean) => {
            setIsOpen(newOpen);
        };

        return (
            <ReasoningContext.Provider
                value={{ isStreaming, isOpen, setIsOpen, duration }}
            >
                <Collapsible
                    className={cn('not-prose mb-4', className)}
                    onOpenChange={handleOpenChange}
                    open={isOpen}
                    {...props}
                >
                    {children}
                </Collapsible>
            </ReasoningContext.Provider>
        );
    }
);

export type ReasoningTriggerProps = ComponentProps<
    typeof CollapsibleTrigger
> & {
    title?: string;
};

export const ReasoningTrigger = memo(
    ({
         className,
         title = 'Reasoning',
         children,
         ...props
     }: ReasoningTriggerProps) => {
        const { isStreaming, isOpen, duration } = useReasoning();

        return (
            <CollapsibleTrigger
                className={cn(
                    'flex items-center gap-2 text-muted-foreground text-sm',
                    className
                )}
                {...props}
            >
                {children ?? (
                    <>
                        <BrainIcon className="size-4" />
                        {isStreaming || duration === 0 ? (
                            <p>Thinking...</p>
                        ) : (
                            <p>Thought for {duration} seconds</p>
                        )}
                    </>
                )}
                <ChevronDownIcon
                    className={cn(
                        'size-4 text-muted-foreground transition-transform',
                        isOpen ? 'rotate-180' : 'rotate-0'
                    )}
                />
            </CollapsibleTrigger>
        );
    }
);

export type ReasoningContentProps = ComponentProps<
    typeof CollapsibleContent
> & {
    children: string;
};

export const ReasoningContent = memo(
    ({ className, children, ...props }: ReasoningContentProps) => (
        <CollapsibleContent
            className={cn(
                'mt-4 text-sm',
                'data-[state=closed]:fade-out-0 data-[state=closed]:slide-out-to-top-2 data-[state=open]:slide-in-from-top-2 text-popover-foreground outline-none data-[state=closed]:animate-out data-[state=open]:animate-in',
                className
            )}
            {...props}
        >
            <Response className="grid gap-2">{children}</Response>
        </CollapsibleContent>
    )
);

Reasoning.displayName = 'Reasoning';
ReasoningTrigger.displayName = 'ReasoningTrigger';
ReasoningContent.displayName = 'ReasoningContent';
